import ssl
import gevent
from gevent.pool import Group
from locust import task
import time
import websocket
from locust import User


class SocketIOUser(User):
    """
    A User that includes a socket io websocket connection.
    You could easily use this a template for plain WebSockets,
    socket.io just happens to be my use case. You can use multiple
    inheritance to combine this with an HttpUser
    (class MyUser(HttpUser, SocketIOUser)
    """

    abstract = True  # 为True标识为非执行类
    greenlet = Group()

    def connect(self, host: str, header=[], **kwargs):
        self.state = True
        self.hosts = host
        # 建立连接步骤，根据实际情况自定义
        try:
            time1 = time.time()
            self.ws = websocket.create_connection(host, header=header, **kwargs)
            time2 = time.time()

            # 连接成功，则统计成功情况
            self.environment.events.request.fire(
                request_type="Connect",
                name=host[0:20],
                response_time=int((time2 - time1) * 1000),
                response_length=0,
                exception=None,
                context=None,
            )
        except Exception as e:
            # 连接失败，则统计失败情况
            self.environment.events.request.fire(
                request_type="Connect",
                name=host[0:20],
                response_time=0,
                response_length=0,
                exception=f"{e}: on {self.hosts}",
                context=f"{e}: on {self.hosts}",
            )
            return

    def spawn(self, send_message_loop=False, send_userList_loop=False):
        # 将自定义处理内容加入协程并开始执行，可根据业务情况自定义增加发送的消息，一般情况下心跳和接收消息为必须
        self.greenlet.add(gevent.spawn(self.sleep_with_heartbeat))
        self.greenlet.add(gevent.spawn(self.receive_loop))
        if send_message_loop:
            self.greenlet.add(gevent.spawn(self.send_message_loop))
        if send_userList_loop:
            self.greenlet.add(gevent.spawn(self.send_user_list_loop))
        self.greenlet.join()

    def send_user_list_loop(self, seconds=2):
        # 自定义发送消息协程，注意协程里sleep要用协程的sleep方法，另外，while循环中，缺少sleep可能导致拥塞，通过gevent.sleep将执行权交出去
        self._before = time.time()
        while self.state:
            if time.time() - self._before >= seconds:
                self.send('{"method":"userList"}')
                self._before = time.time()
            gevent.sleep(0.01)

    def send_message_loop(self, seconds=1):
        # 自定义发送消息协程，这里间隔1秒发送一个消息
        self._before2 = time.time()
        while self.state:
            if time.time() - self._before2 >= seconds:
                self.send('{"method":"message","data":"123"}')
                self._before2 = time.time()
            gevent.sleep(0.01)

    def on_message(self, message):  # override this method in your subclass for custom handling
        # 根据接受到的消息内容进行分类统计，注意根据修改if else里面的判断，也可以自己写正则等方法进行处理
        if str(message) == "3":
            name = 'heartbeat'
        elif 'userList' in str(message):
            name = 'userList'
        elif 'message' in str(message):
            name = 'message'
            # 如果按业务约定，收到服务器消息需要回复，在此处理即可, 以下为示例
            if 'needResponse' in str(message):
                self.send('this is a response message.')
        else:
            # print(f"Received unexpected message: {message}")
            return
        self.environment.events.request.fire(
            request_type="RSV",
            name=name,
            response_time=0,
            response_length=len(str(message)),
            exception=None,
            context=None,
        )

    def receive_loop(self):
        # 接收消息处理协程
        while self.state:
            try:
                message = self.ws.recv()
                self.on_message(message)
                gevent.sleep(0.01)
            except Exception as e:
                self.environment.events.request.fire(
                    request_type="RSV",
                    name="RSV Error",
                    response_time=0,
                    response_length=0,
                    exception=f"{e} on {self.hosts}",
                    context=f"{e} on {self.hosts}",
                )
                if 'closed' in str(e):
                    self.state = 0
                return

    def send(self, body, context={}):
        # print(str(body))
        # 根据发送消息内容进行分类统计，并发送消息，注意根据修改if else里面的判断，也可以自己写正则等方法进行处理
        if str(body) == "2":
            name = 'heartbeat'
        elif 'userList' in str(body):
            name = 'userList'
        elif 'message' in str(body):
            name = 'message'
        else:
            # print(f"Send unexpected message: {message}")
            return

        # request_type 可根据实际情况修改
        # name 可根据实际情况修改
        # response_time 可自己想办法计算，但是异步收发计算难度较大，无法举例
        try:
            self.ws.send(body)
            self.environment.events.request.fire(
                request_type="SEND",
                name=name,
                response_time=None,
                response_length=len(str(body)),
                exception=None,
                context=None,
            )
        except Exception as e:
            self.environment.events.request.fire(
                request_type="SEND",
                name=name,
                response_time=0,
                response_length=0,
                exception=f"{e} on host {self.hosts}",
                context=f"{e} on host {self.hosts}",
            )
            return

    def sleep_with_heartbeat(self, seconds=5):
        # 心跳维持协程，同时部分定时发送消息可在心跳中进行同步发送
        self._before3 = time.time()
        while self.state:
            if time.time() - self._before3 >= seconds:
                self.send('2')
                self._before3 = time.time()
            gevent.sleep(0.01)

    def leave(self):
        # 发送41消息退出房间，使用此方法，则不能使用spawn方法
        self.send('41')


class MySocketIOUser1(SocketIOUser):
    # 这里是用例实例类，可写多个task，不同的流程进行混合调用, 多个类情况下可通过task(数字)进行控制比例，当然也可以写多个类
    # 注：因为websocket是长连接，在释放前不会循环调用

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        # url如需其他参数则根据实际情况进行处理
        self.url = f'ws://localhost:8765'
        # 其他需要初始化的内容

    @task(1)
    def my_task1(self):
        # 用例
        self.connect(
            host=f'{self.url}',
            sslopt={"cert_reqs": ssl.CERT_NONE})
        self.spawn(send_message_loop=True, send_userList_loop=True)


# 调试： locust -f locust_websocket.py -u 20 -r 5 -t 60 --headless
# web界面主节点 locust -f locust_websocket.py --master
# web界面子节点 locust -f locust_websocket.py --worker

# 无界面服务器分布式压测，到预期的worker数后开始压测
# locust -f locust_websocket.py --master --web-host="127.0.0.1" --headless --expect-workers=120 -u 10000 -r 20 -t 30m --html=../perf.html
# nohup locust -f locust_websocket.py --worker --master-host=xxx.xxx.xxx.xxx > /dev/null 2>&1 &
